import string
import numpy as np
import cv2

n = 0    #定义鼠标按下的次数
ix = 0   # x,y 坐标的临时存储
iy = 0
rect = (0,0,0,0) #前景区域
def draw_rectangle(event,x,y,flags,param):
    global n,ix,iy,rect
    if event==cv2.EVENT_LBUTTONDOWN :
        if n == 0:    #首次按下保存坐标值
            n+=1
            ix,iy = x,y
            print(x,y)
        else:        #第二次按下显示矩形
            n+=1
            rect = (ix,iy,(x-ix),(y-iy))#前景区域
            print(x,y)

def add_background(path:string,backgroundPath:string,outPath:string,name:string):
    #读取图像
    img = cv2.imread(path+name+".jpg")
    bg=cv2.imread(backgroundPath+name+"_background.jpg")
    bg=cv2.resize(bg,(img.shape[1],img.shape[0]))
    mask = np.zeros(img.shape[:2],np.uint8)
    bgdModel = np.zeros((1,65),np.float64)
    fgdModel = np.zeros((1,65),np.float64)
    #选择区域 左上到右下矩形
    rect=(5,5,img.shape[1]-10,img.shape[0]-5)
    cv2.namedWindow("img")
    #前景提取
    # cv2.setMouseCallback("img",draw_rectangle)#绑定鼠标
    # while(n != 2):
    #     cv2.imshow("img",img)
    #     cv2.waitKey(2)
    cv2.grabCut(img,mask,rect,bgdModel,fgdModel,20,cv2.GC_INIT_WITH_RECT)
    mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
    img = img*mask2[:,:,np.newaxis]
    mask3=np.where((mask2==1),0,1).astype('uint8')
    bg=bg*mask3[:,:,np.newaxis]
    #显示图像
    cv2.imshow("img",bg+img)
    cv2.imwrite(outPath+"front/"+name+".jpg",img)
    cv2.imwrite(outPath+name+".jpg",bg+img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__=="__main__":
    path="./graph/"
    backgroundPath="./background/"
    outPath="./add_background/"
    for i in range(1,6):
        add_background(path,backgroundPath,outPath,"0"+str(i))
    # add_background(path,backgroundPath,outPath,"01")